{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3219fa3f",
   "metadata": {},
   "source": [
    "# 目标钩子\n",
    "\n",
    "此文件包含TVM框架中目标钩子功能的单元测试，主要测试外部代码生成相关的功能。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b1b62214",
   "metadata": {
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "import set_env"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e06abb2e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import numpy as np\n",
    "import logging\n",
    "from pathlib import Path\n",
    "\n",
    "import tvm\n",
    "import tvm.testing\n",
    "from tvm import relay, IRModule\n",
    "from tvm.contrib import utils\n",
    "from tvm.relay.op.annotation import compiler_begin, compiler_end\n",
    "from tvm import relay, runtime, testing\n",
    "logging.basicConfig(level=logging.INFO)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d37bae4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def set_external_func_attr(func, compiler, ext_symbol):\n",
    "    func = func.with_attr(\"Primitive\", tvm.tir.IntImm(\"int32\", 1))\n",
    "    func = func.with_attr(\"Compiler\", compiler)\n",
    "    func = func.with_attr(\"global_symbol\", ext_symbol)\n",
    "    return func\n",
    "\n",
    "def update_lib(lib, source_dir):\n",
    "    source_dir = Path(source_dir)\n",
    "    contrib_path = source_dir/\"src/runtime/contrib\"\n",
    "\n",
    "    kwargs = {}\n",
    "    kwargs[\"options\"] = [\"-O2\", \"-std=c++17\", f\"-I{contrib_path}\"]\n",
    "    tmp_path = utils.tempdir()\n",
    "    lib_name = \"lib.so\"\n",
    "    lib_path = tmp_path.relpath(lib_name)\n",
    "    lib.export_library(lib_path, fcompile=False, **kwargs)\n",
    "    lib = tvm.runtime.load_module(lib_path)\n",
    "    return lib\n",
    "\n",
    "def check_result(\n",
    "    mod, map_inputs, out_shape, result, tol=1e-5, \n",
    "    target=\"llvm\", device=tvm.cpu(), \n",
    "    source_dir=\"/media/pc/data/board/arria10/lxw/tasks/tvm\"):\n",
    "    with tvm.transform.PassContext(opt_level=3, disabled_pass=[\"AlterOpLayout\"]):\n",
    "        exe = relay.vm.compile(mod, target=target)\n",
    "    code, lib = exe.save()\n",
    "    lib = update_lib(lib, source_dir=source_dir)\n",
    "    exe = runtime.vm.Executable.load_exec(code, lib)\n",
    "    vm = runtime.vm.VirtualMachine(exe, device)\n",
    "    out = vm.run(**map_inputs)\n",
    "    tvm.testing.assert_allclose(out.numpy(), result, rtol=tol, atol=tol)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "344ee6cb",
   "metadata": {},
   "source": [
    "## 测试内联模式下不使用目标实例的TIR外部代码生成"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "989379fb",
   "metadata": {},
   "source": [
    "此测试用例验证当不使用具体目标实例时，内联模式的TIR外部代码生成是否正常工作。\n",
    "\n",
    "它将简单的加法算子标记为外部函数，该函数在自定义代码生成时会被替换为减法算子。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "eb5e4505",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #A2F\">@main</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x: Tensor[(<span style=\"color: #008000\">8</span>), float32], <span style=\"color: #A2F; font-weight: bold\">%</span>y: Tensor[(<span style=\"color: #008000\">8</span>), float32]) {\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #A2F; font-weight: bold\">=</span> fn (<span style=\"color: #A2F; font-weight: bold\">%</span>x0: Tensor[(<span style=\"color: #008000\">8</span>), float32], <span style=\"color: #A2F; font-weight: bold\">%</span>y0: Tensor[(<span style=\"color: #008000\">8</span>), float32], Primitive<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>, Compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;example_target_hook&quot;</span>, global_symbol<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;replace_add_with_subtract&quot;</span>) {\n",
       "    add(<span style=\"color: #A2F; font-weight: bold\">%</span>x0, <span style=\"color: #A2F; font-weight: bold\">%</span>y0)\n",
       "  };\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x, <span style=\"color: #A2F; font-weight: bold\">%</span>y)\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# 定义输入数据形状\n",
    "shape = (8,)\n",
    "# 生成随机测试数据\n",
    "x_data = np.random.randint(255, size=shape).astype(\"float32\")\n",
    "y_data = np.random.randint(255, size=shape).astype(\"float32\")\n",
    "inputs = {\"x\": x_data, \"y\": y_data}\n",
    "\n",
    "# 创建Relay变量\n",
    "x0 = relay.var(\"x0\", shape=shape, dtype=\"float32\")\n",
    "y0 = relay.var(\"y0\", shape=shape, dtype=\"float32\")\n",
    "# 定义简单的加法操作\n",
    "z = x0 + y0\n",
    "# 创建函数\n",
    "f = relay.Function([x0, y0], z)\n",
    "# 设置外部函数属性，指定目标钩子和全局符号名称\n",
    "f = set_external_func_attr(f, \"example_target_hook\", \"replace_add_with_subtract\")\n",
    "\n",
    "# 创建主函数的输入变量\n",
    "x = relay.var(\"x\", shape=(8,), dtype=\"float32\")\n",
    "y = relay.var(\"y\", shape=(8,), dtype=\"float32\")\n",
    "# 调用外部函数\n",
    "call = relay.Call(f, [x, y])\n",
    "# 创建IR模块\n",
    "func = IRModule.from_expr(call)\n",
    "\n",
    "# 验证结果，预期输出为x_data - y_data（因为加法被替换为减法）\n",
    "check_result(func, inputs, (8,), x_data - y_data)\n",
    "func.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "efc21e41",
   "metadata": {},
   "source": [
    "## 测试带有目标实例的轮廓模式下的TIR外部代码生成"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c26bd0f7",
   "metadata": {},
   "source": [
    "此测试用例验证当使用具体目标实例（包括自定义属性）时，轮廓模式的TIR外部代码生成是否正常工作。\n",
    "    \n",
    "它演示了如何将目标属性传递到自定义pass中。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "6755fd3f",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:te_compiler:Using injective.cpu for multiply based on highest priority (10)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #A2F\">@main</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x: Tensor[(<span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>), float32] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">4</span>:<span style=\"color: #008000\">40</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>y: Tensor[(<span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>), float32] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">4</span>:<span style=\"color: #008000\">44</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">8</span>), float32] {\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #A2F; font-weight: bold\">=</span> <span style=\"color: #A2F\">@replace_add_with_subtract</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x, <span style=\"color: #A2F; font-weight: bold\">%</span>y) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>), float32] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">4</span>:<span style=\"color: #008000\">13</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>;\n",
       "  multiply(<span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">2</span>f <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>float32 span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">4</span>:<span style=\"color: #008000\">54</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>), float32] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">4</span>:<span style=\"color: #008000\">13</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "}\n",
       "\n",
       "<span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #A2F\">@replace_add_with_subtract</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x1: Tensor[(<span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>), float32] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">12</span>:<span style=\"color: #008000\">13</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, <span style=\"color: #A2F; font-weight: bold\">%</span>y1: Tensor[(<span style=\"color: #008000\">8</span>), float32] <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>), float32] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">12</span>:<span style=\"color: #008000\">18</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>, Inline<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>, Primitive<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>, Compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;example_target_hook&quot;</span>, global_symbol<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;replace_add_with_subtract&quot;</span>) <span style=\"color: #A2F; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">8</span>), float32] {\n",
       "  add(<span style=\"color: #A2F; font-weight: bold\">%</span>x1, <span style=\"color: #A2F; font-weight: bold\">%</span>y1) <span style=\"color: #A2F; font-weight: bold\">/*</span> ty<span style=\"color: #A2F; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">8</span>), float32] span<span style=\"color: #A2F; font-weight: bold\">=</span>from_string:<span style=\"color: #008000\">12</span>:<span style=\"color: #008000\">13</span> <span style=\"color: #A2F; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# 定义输入数据形状\n",
    "shape = (8,)\n",
    "# 生成随机测试数据\n",
    "x_data = np.random.randint(255, size=shape).astype(\"float32\")\n",
    "y_data = np.random.randint(255, size=shape).astype(\"float32\")\n",
    "inputs = {\"x\": x_data, \"y\": y_data}\n",
    "# 使用带钩子的目标实例进行编译，以演示将目标属性传递到自定义pass中\n",
    "host_target = tvm.target.Target(\"llvm\")\n",
    "generic_target = tvm.target.Target(\"llvm\", host=host_target)\n",
    "# 创建带有自定义属性的外部代码生成目标\n",
    "extern_codegen_target = tvm.target.Target(\n",
    "    \"example_target_hook -example_attribute=42\", host=host_target\n",
    ")\n",
    "# 从文本创建IR模块\n",
    "mod = tvm.relay.fromtext(\n",
    "    \"\"\"\n",
    "        #[version = \"0.0.5\"]\n",
    "        def @main(%x: Tensor[(8), float32], %y: Tensor[(8), float32]) -> Tensor[(8), float32] {\n",
    "            @replace_add_with_subtract(%x, %y) * 2.0f\n",
    "        }\n",
    "\n",
    "        def @replace_add_with_subtract(%x: Tensor[(8), float32], %y: Tensor[(8), float32],\n",
    "                                        Inline=1,\n",
    "                                        Primitive=1,\n",
    "                                        Compiler=\"example_target_hook\",\n",
    "                                        global_symbol=\"replace_add_with_subtract\") -> Tensor[(8), float32] {\n",
    "            %x + %y  // 将会被自定义pass重写为实现%x - %y - 42.0f的TIR\n",
    "        }\n",
    "    \"\"\"\n",
    ")\n",
    "\n",
    "# 验证结果，预期输出为(x_data - y_data - 42.0) * 2.0\n",
    "check_result(\n",
    "    mod,\n",
    "    inputs,\n",
    "    (8,),\n",
    "    (x_data - y_data - 42.0) * 2.0,\n",
    "    target=[generic_target, extern_codegen_target],\n",
    ")\n",
    "mod.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f09cc66",
   "metadata": {},
   "source": [
    "## 测试运行时模块生成"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0a106561",
   "metadata": {},
   "source": [
    "此测试用例验证带有 `'tir_to_runtime'` 属性的函数是否能正确触发 TIR 到运行时的代码生成流程。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "9b1ebbaa",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #A2F\">@main</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x: Tensor[(<span style=\"color: #008000\">8</span>), float32], <span style=\"color: #A2F; font-weight: bold\">%</span>y: Tensor[(<span style=\"color: #008000\">8</span>), float32]) {\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #A2F; font-weight: bold\">=</span> fn (<span style=\"color: #A2F; font-weight: bold\">%</span>x0: Tensor[(<span style=\"color: #008000\">8</span>), float32], <span style=\"color: #A2F; font-weight: bold\">%</span>y0: Tensor[(<span style=\"color: #008000\">8</span>), float32], Primitive<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>, Compiler<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;example_target_hook&quot;</span>, global_symbol<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;replace_add_with_subtract&quot;</span>, tir_to_runtime<span style=\"color: #A2F; font-weight: bold\">=</span><span style=\"color: #008000; font-weight: bold\">True</span>) {\n",
       "    add(<span style=\"color: #A2F; font-weight: bold\">%</span>x0, <span style=\"color: #A2F; font-weight: bold\">%</span>y0)\n",
       "  };\n",
       "  <span style=\"color: #A2F; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>(<span style=\"color: #A2F; font-weight: bold\">%</span>x, <span style=\"color: #A2F; font-weight: bold\">%</span>y)\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# 定义输入数据形状\n",
    "shape = (8,)\n",
    "# 生成随机测试数据\n",
    "x_data = np.random.randint(255, size=shape).astype(\"float32\")\n",
    "y_data = np.random.randint(255, size=shape).astype(\"float32\")\n",
    "inputs = {\"x\": x_data, \"y\": y_data}\n",
    "\n",
    "# 创建Relay变量\n",
    "x0 = relay.var(\"x0\", shape=shape, dtype=\"float32\")\n",
    "y0 = relay.var(\"y0\", shape=shape, dtype=\"float32\")\n",
    "# 定义简单的加法操作\n",
    "z = x0 + y0\n",
    "# 创建函数\n",
    "func = relay.Function([x0, y0], z)\n",
    "# 设置外部函数属性\n",
    "func = set_external_func_attr(func, \"example_target_hook\", \"replace_add_with_subtract\")\n",
    "# 添加触发TIRToRuntime代码生成的钩子\n",
    "func = func.with_attr(\"tir_to_runtime\", True)\n",
    "\n",
    "# 创建主函数的输入变量\n",
    "x = relay.var(\"x\", shape=(8,), dtype=\"float32\")\n",
    "y = relay.var(\"y\", shape=(8,), dtype=\"float32\")\n",
    "# 调用外部函数\n",
    "call = relay.Call(func, [x, y])\n",
    "# 创建IR模块\n",
    "func = IRModule.from_expr(call)\n",
    "\n",
    "# 验证结果，预期输出为x_data * y_data\n",
    "check_result(func, inputs, (8,), x_data * y_data)\n",
    "func.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f3d0da4",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "aix",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
